package com.da.javatest.thread.task;


import com.google.common.collect.Lists;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Vector;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @description:
 * @author: yajun.dyj
 * @create: 2018-12-24 19:35
 **/
public class MultiThreadTaskRunner {

    /**
     * 多线程执行任务，用户需提供处理单个任务的方法 singleRunner
     * 返回一个list，list中的每个元素为一个pair，A为入参，B为结果
     *
     * @param taskList     任务列表
     * @param singleRunner 单个任务处理方法
     * @param threadPool   线程池
     * @param threadNum    线程数量
     * @param <T>          任务类型
     * @param <R>          任务返回值
     * @return
     */
    public static <T, R> List<Pair<T, R>> runTasks(List<T> taskList, Function<T, R> singleRunner,
                                                   ExecutorService threadPool, int threadNum) {
        return runTasksListRunner(taskList,
            (mtTaskList) -> mtTaskList.stream()
                .map(task -> new Pair<>(task, singleRunner.apply(task))).collect(Collectors.toList()),
            threadPool, threadNum);
    }

    /**
     * 多线程执行任务，用户需提供处理批量任务的方法 listRunner
     * 返回一个list，list中的每个元素为一个pair，A为入参，B为结果
     *
     * @param taskList   待处理的任务列表
     * @param listRunner 批量任务处理方法
     * @param threadPool 线程池
     * @param threadNum  线程个数
     * @param <T>        待处理的任务类型
     * @param <R>        返回值类型
     * @return
     */
    public static <T, R> List<Pair<T, R>> runTasksListRunner(List<T> taskList,
                                                             Function<List<T>, List<Pair<T, R>>> listRunner,
                                                             ExecutorService threadPool, int threadNum) {
        if (CollectionUtils.isEmpty(taskList) || threadNum < 1) {
            return Lists.newArrayList();
        }
        List<List<T>> mtTaskLists = Lists.newArrayList();
        //通过轮转的方式分配任务
        buildMtTaskLists(taskList, threadNum, mtTaskLists);
        CountDownLatch countDownLatch = new CountDownLatch(mtTaskLists.size());
        List<Pair<T, R>> resList = new Vector<>();
        mtTaskLists.forEach((mtTaskList) -> threadPool.submit(() -> {
            try {
                ExceptionHandler.ignore(() -> resList.addAll(listRunner.apply(mtTaskList)));
            } finally {
                countDownLatch.countDown();
            }
        }));
        ExceptionHandler.ignore(countDownLatch::await);
        return resList;
    }

    private static <T> void buildMtTaskLists(List<T> taskList, int threadNum, List<List<T>> mtTaskLists) {
        int[] index = {0};
        taskList.forEach(task -> {
            if (mtTaskLists.size() == index[0]) {
                mtTaskLists.add(Lists.newArrayList(task));
            } else {
                mtTaskLists.get(index[0]).add(task);
            }
            index[0] = (index[0] + 1) % threadNum;
        });
    }

}
